package com.cyy.chat.model;

import com.cyy.chat.model.result.EmbeddingObj;
import com.cyy.chat.model.result.EmbeddingsApiResult;
import com.google.gson.Gson;
import com.google.gson.reflect.TypeToken;
import jakarta.validation.Valid;
import okhttp3.*;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import java.net.InetSocketAddress;
import java.net.Proxy;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.TimeUnit;


@Component
public class EmbeddingModel {

    @Value("${huggingface.url}")
    private String url;
    @Value("${huggingface.token}")
    private String token;

    public EmbeddingsApiResult doEmbeddingByHuggingface(String msg) {
        String json = "{\"inputs\": \"" + msg + "\"}";
        RequestBody requestBody = RequestBody.create(json, MediaType.get("application/json; charset=utf-8"));

        Request request = new Request.Builder()
                .url(url)
                .addHeader("Authorization", "Bearer " + token)
                .addHeader("Content-Type", "application/json")
                .post(requestBody)
                .build();
        Proxy proxy = new Proxy(Proxy.Type.SOCKS, new InetSocketAddress("127.0.0.1", 7890));

        // 直接创建 OkHttpClient 实例
        OkHttpClient client = new OkHttpClient.Builder()
                .proxy(proxy)
                .connectTimeout(30, TimeUnit.SECONDS)
                .readTimeout(30, TimeUnit.SECONDS)
                .writeTimeout(30, TimeUnit.SECONDS)
                .build();

        try (Response response = client.newCall(request).execute()) {
            if (response.code() == 200 && response.body() != null) {
                String responseStr = response.body().string();

                Gson gson = new Gson();
                List<List<Float>> embeddingArray = gson.fromJson(responseStr, new TypeToken<List<List<Float>>>() {}.getType());

                EmbeddingObj obj = new EmbeddingObj();
                obj.setEmbedding(embeddingArray.get(0));
                obj.setIndex(0);

                EmbeddingsApiResult result = new EmbeddingsApiResult();
                result.setModel("BAAI/bge-small-zh-v1.5");
                result.setObject("embedding");
                result.setUsage(null);
                result.setData(Arrays.asList(obj));

                return result;
            } else {
                System.err.println("HuggingFace embedding error: " + response.code() + " - " + (response.body() != null ? response.body().string() : "No response body"));
                return null;
            }
        } catch (Exception e) {
            e.printStackTrace();
            return null;
        }
    }



}
